#include "cx_temp.h"
#include <sddl.h> // ConvertSidToStringSidA
#include <array>
#include "nlohmann/json.hpp"

namespace CXLib {

	using namespace nlohmann;

	/// <summary>
	/// Format process token Info
	/// </summary>
	/// <param name="token"></param>
	/// <returns></returns>
	std::string cxFormatProcessTokenInfo(const ProcessTokenInfo& token)
	{
		json jResult;
		jResult["sessionId"] = token.sessionId;
		jResult["tokenIsElevated"] = token.tokenIsElevated;
		jResult["integrityLevel"] = token.integrityLevel;
		jResult["owner"] = {
			{"sid", token.owner.sid},
			{"name", token.owner.name},
			{"domain", token.owner.domain}
		};
		jResult["user"] = {
			{"sid", token.user.sid},
			{"name", token.user.name},
			{"domain", token.user.domain}
		};
		jResult["primaryGroup"] = {
			{"sid", token.primaryGroup.sid},
			{"name", token.primaryGroup.name},
			{"domain", token.primaryGroup.domain}
		};
		for (const auto& group : token.groups) {
			jResult["groups"].push_back({
				{"sid", group.sid},
				{"name", group.name},
				{"domain", group.domain},
				{"attr", group.attr},
				{"attrDesc", group.attrDesc}
			});
		}
		for (const auto& priv : token.privileges) {
			jResult["privileges"].push_back({
				{"name", priv.name},
				{"state", priv.state}
			});
		}
		for (const auto& ace : token.daclAces) {
			jResult["daclAces"].push_back({
				{"type", ace.type},
				{"flags", ace.flags},
				{"rights", ace.rights},
				{"object_guid", ace.object_guid},
				{"inherit_object_guid", ace.inherit_object_guid},
				{"account_sid", ace.account_sid},
				{"desc", {
					{"type", ace.desc.type_desc},
					{"flags", ace.desc.flags_desc},
					{"rights", ace.desc.rights_desc}
				}}
			});
		}


		return jResult.dump();
	}

	/// <summary>
	/// Change privilege of process token
	/// </summary>
	/// <param name="hToken"></param>
	/// <param name="lpszPrivilege"></param>
	/// <param name="bEnablePrivilege"></param>
	/// <returns></returns>
	DWORD cxSetPrivilege(
		HANDLE hToken,          // access token handle
		LPCTSTR lpszPrivilege,  // name of privilege to enable/disable
		BOOL bEnablePrivilege   // to enable or disable privilege
	)
	{
		if (!hToken)
		{
			return ERROR_INVALID_HANDLE;
		}

		TOKEN_PRIVILEGES tp;
		LUID luid;

		if (!LookupPrivilegeValue(
			NULL,            // lookup privilege on local system
			lpszPrivilege,   // privilege to lookup 
			&luid))        // receives LUID of privilege
		{
			return ::GetLastError();
		}

		tp.PrivilegeCount = 1;
		tp.Privileges[0].Luid = luid;
		if (bEnablePrivilege)
			tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
		else
			tp.Privileges[0].Attributes = 0;

		// Enable the privilege or disable all privileges.

		if (!AdjustTokenPrivileges(
			hToken,
			FALSE,
			&tp,
			sizeof(TOKEN_PRIVILEGES),
			(PTOKEN_PRIVILEGES)NULL,
			(PDWORD)NULL))
		{
			return GetLastError();
		}

		if (GetLastError() == ERROR_NOT_ALL_ASSIGNED)
		{
			return GetLastError();
		}

		return ERROR_SUCCESS;
	}


	/// <summary>
	/// Get privileges of process token
	/// </summary>
	/// <param name="hToken"></param>
	/// <param name="privileges"></param>
	/// <returns></returns>
	DWORD cxGetTokenPrivileges(HANDLE hToken, std::vector<PrivilegeInfo>& privileges)
	{
		if (!hToken)
		{
			return ERROR_INVALID_HANDLE;
		}

		DWORD len = 0;
		GetTokenInformation(hToken, TokenPrivileges, nullptr, 0, &len);
		if (len == 0) {
			return ERROR_INVALID_PARAMETER;
		}

		auto buffer = (BYTE*)malloc(len);
		if (!buffer) return ERROR_OUTOFMEMORY;
		if (GetTokenInformation(hToken, TokenPrivileges, buffer, len, &len))
		{

			auto pPrivs = (PTOKEN_PRIVILEGES)buffer;

			for (DWORD i = 0; i < pPrivs->PrivilegeCount; ++i)
			{
				LUID_AND_ATTRIBUTES &laa = pPrivs->Privileges[i];

				std::array<char, 256> name;
				DWORD cchName = 256;
				PrivilegeInfo info;
				if (LookupPrivilegeNameA(NULL, &laa.Luid, name.data(), &cchName))
				{
					info.name.assign(name.data(), cchName);
					if (laa.Attributes & SE_PRIVILEGE_ENABLED)
						info.state = "SE_PRIVILEGE_ENABLED";
					else if (laa.Attributes & SE_PRIVILEGE_ENABLED_BY_DEFAULT)
						info.state = "SE_PRIVILEGE_ENABLED_BY_DEFAULT";
					else if (laa.Attributes & SE_PRIVILEGE_REMOVED)
						info.state = "SE_PRIVILEGE_REMOVED";
					else
						info.state = "SE_PRIVILEGE_DISABLED";
					privileges.push_back(info);
				}
			}
		}

		free(buffer);

		return ERROR_SUCCESS;
	}

	/**
  * @brief Get the token information of the specified process
  * @param hProcess hProcess (requires TOKEN-QRY permission)
  * @param token output parameters, save token related information
  * @return 0 indicates success, while others are Win32 error codes
  * @note only supports Windows, the caller needs to ensure that hProcess is valid
  */
	DWORD cxGetProcessTokenInfo(HANDLE hProcess, ProcessTokenInfo& token)
	{
		if (!hProcess)
		{
			return ERROR_INVALID_HANDLE;
		}

		HANDLE hToken;
		if (!OpenProcessToken(hProcess, TOKEN_QUERY, &hToken)) return ::GetLastError();

		struct TokenInfoType {
			TOKEN_INFORMATION_CLASS type;
			const char* name;
		};

		TokenInfoType tokenTypes[] = {
			{ TokenUser,                   "TokenUser" },
			{ TokenGroups,                 "TokenGroups" },
			{ TokenPrivileges,             "TokenPrivileges" },
			{ TokenOwner,                  "TokenOwner" },
			{ TokenPrimaryGroup,           "TokenPrimaryGroup" },
			{ TokenDefaultDacl,            "TokenDefaultDacl" },
			{ TokenSessionId,              "TokenSessionId" },
			{ TokenIntegrityLevel,         "TokenIntegrityLevel" },
			{ TokenElevation,              "TokenElevation" }
		};

		for (const auto& info : tokenTypes) {
			DWORD len = 0;
			GetTokenInformation(hToken, info.type, nullptr, 0, &len);
			if (len == 0) {
				continue;
			}

			auto buffer = (BYTE*)malloc(len);
			if (!buffer) continue;
			if (GetTokenInformation(hToken, info.type, buffer, len, &len)) {
				
				if (info.type == TokenUser) {
					auto pUser = (PTOKEN_USER)buffer;
		
					if (char* sidStr = nullptr; ConvertSidToStringSidA(pUser->User.Sid, &sidStr)) {
						token.user.sid = sidStr;
						LocalFree(sidStr); 
					}

					std::array<char, 256> name;
					std::array<char, 256> domain;
					DWORD cchName = 256;
					DWORD cchDomain = 256;
					SID_NAME_USE sidType;
					if (LookupAccountSidA(NULL, pUser->User.Sid, name.data(), &cchName, domain.data(), &cchDomain, &sidType)) {
						token.user.name.assign(name.data(), cchName);
						token.user.domain.assign(domain.data(), cchDomain);
					}
				}
				else if (info.type == TokenGroups) {
					auto pGroups = (PTOKEN_GROUPS)buffer;

					for (DWORD i = 0; i < pGroups->GroupCount; ++i) {
						SidInfo sidInfo;
						if (char* sidStr = nullptr; ConvertSidToStringSidA(pGroups->Groups[i].Sid, &sidStr)) {
							sidInfo.sid = sidStr;
							LocalFree(sidStr);
						}

						std::array<char, 256> name;
						std::array<char, 256> domain;
						DWORD cchName = 256;
						DWORD cchDomain = 256;
						if (SID_NAME_USE sidType; LookupAccountSidA(NULL, pGroups->Groups[i].Sid, name.data(), &cchName, domain.data(), &cchDomain, &sidType)) {
							sidInfo.name.assign(name.data(), cchName);
							sidInfo.domain.assign(domain.data(), cchDomain);
						}

						sidInfo.attr = pGroups->Groups[i].Attributes;
						std::vector<std::string> attrList;
						if (sidInfo.attr & SE_GROUP_ENABLED)
							attrList.emplace_back("SE_GROUP_ENABLED");
						if (sidInfo.attr & SE_GROUP_ENABLED_BY_DEFAULT)
							attrList.emplace_back("SE_GROUP_ENABLED_BY_DEFAULT");
						if (sidInfo.attr & SE_GROUP_LOGON_ID)
							attrList.emplace_back("SE_GROUP_LOGON_ID");
						if (sidInfo.attr & SE_GROUP_OWNER)
							attrList.emplace_back("SE_GROUP_OWNER");
						if (sidInfo.attr & SE_GROUP_USE_FOR_DENY_ONLY)
							attrList.emplace_back("SE_GROUP_USE_FOR_DENY_ONLY");
						if (sidInfo.attr & SE_GROUP_INTEGRITY)
							attrList.emplace_back("SE_GROUP_INTEGRITY");
						if (sidInfo.attr & SE_GROUP_INTEGRITY_ENABLED)
							attrList.emplace_back("SE_GROUP_INTEGRITY_ENABLED");
						if (sidInfo.attr & SE_GROUP_RESOURCE)
							attrList.emplace_back("SE_GROUP_RESOURCE");
						if (sidInfo.attr & SE_GROUP_MANDATORY)
							attrList.emplace_back("SE_GROUP_MANDATORY");
						if (!attrList.empty()) {
							for (size_t j = 0; j < attrList.size(); ++j) {
								sidInfo.attrDesc += attrList[j];
								if (j != attrList.size() - 1)
									sidInfo.attrDesc += " | ";
							}
						}
						token.groups.emplace_back(std::move(sidInfo));
					}
				}
				else if (info.type == TokenPrimaryGroup) {
					auto pPrimary = (PTOKEN_PRIMARY_GROUP)buffer;

					if (char* sidStr = nullptr; ConvertSidToStringSidA(pPrimary->PrimaryGroup, &sidStr)) {
						token.primaryGroup.sid = sidStr;
						LocalFree(sidStr);
					}

					std::array<char, 256> name;
					std::array<char, 256> domain;
					DWORD cchName = 256;
					DWORD cchDomain = 256;
					SID_NAME_USE sidType;
					if (LookupAccountSidA(NULL, pPrimary->PrimaryGroup, name.data(), &cchName, domain.data(), &cchDomain, &sidType)) {
						token.primaryGroup.name.assign(name.data(), cchName);
						token.primaryGroup.domain.assign(domain.data(), cchDomain);
					}
				}
				else if (info.type == TokenOwner) {
					auto pOwner = (PTOKEN_OWNER)buffer;

					if (char* sidStr = nullptr; ConvertSidToStringSidA(pOwner->Owner, &sidStr)) {
						token.owner.sid = sidStr;
						LocalFree(sidStr);
					}

					std::array<char, 256> name;
					std::array<char, 256> domain;
					DWORD cchName = 256;
					DWORD cchDomain = 256;
					SID_NAME_USE sidType;
					if (LookupAccountSidA(NULL, pOwner->Owner, name.data(), &cchName, domain.data(), &cchDomain, &sidType)) {
						token.owner.name.assign(name.data(), cchName);
						token.owner.domain.assign(domain.data(), cchDomain);
					}
				}
				else if (info.type == TokenSessionId) {
					token.sessionId = *(DWORD*)buffer;
				}
				else if (info.type == TokenElevation) {
					auto elev = (TOKEN_ELEVATION*)buffer;
					token.tokenIsElevated = elev->TokenIsElevated;
				}
				else if (info.type == TokenIntegrityLevel) {
					auto pLabel = (PTOKEN_MANDATORY_LABEL)buffer;
					auto dwIntegrityLevel = *GetSidSubAuthority(
						pLabel->Label.Sid,
						(DWORD)(*GetSidSubAuthorityCount(pLabel->Label.Sid) - 1)
					);
					token.integrityLevel = dwIntegrityLevel;
				}
				else if (info.type == TokenDefaultDacl)
				{
					auto pDacl = (PTOKEN_DEFAULT_DACL)buffer;
					if (pDacl->DefaultDacl)
					{
						ACL *acl = pDacl->DefaultDacl;
						USHORT aceCount = acl->AceCount;
						for (USHORT aceIdx = 0; aceIdx < aceCount; ++aceIdx)
						{
							CXSddl::AceInfo aceInfo;
							void *ace = nullptr;
							if (GetAce(acl, aceIdx, &ace))
							{
								switch (auto aceHeader = (ACE_HEADER*)ace; aceHeader->AceType)
								{
								case ACCESS_ALLOWED_ACE_TYPE:
								{
									aceInfo.type = "SDDL_ACCESS_ALLOWED";

									auto allowedAce = (ACCESS_ALLOWED_ACE *)ace;
									auto sid = &allowedAce->SidStart;
									DWORD mask = allowedAce->Mask;
									
									if (char* sidStr = nullptr; ConvertSidToStringSidA(sid, &sidStr)) {
										aceInfo.account_sid = sidStr;
										LocalFree(sidStr);
									}
									
									char maskStr[16];
									sprintf_s(maskStr, "0x%08X", mask);
									aceInfo.rights = maskStr;
									
									break;
								}
								case ACCESS_DENIED_ACE_TYPE:
								{
									aceInfo.type = "SDDL_ACCESS_DENIED";

									auto deniedAce = (ACCESS_DENIED_ACE *)ace;
									auto sid = &deniedAce->SidStart;
									DWORD mask = deniedAce->Mask;
									
									if (char* sidStr = nullptr; ConvertSidToStringSidA(sid, &sidStr)) {
										aceInfo.account_sid = sidStr;
										LocalFree(sidStr);
									}
									
									char maskStr[16];
									sprintf_s(maskStr, "0x%08X", mask);
									aceInfo.rights = maskStr;
									
									break;
								}
								default:
								{
									aceInfo.type = std::to_string(aceHeader->AceType);
								}
									break;
								}
							
								token.daclAces.emplace_back(std::move(aceInfo));
							}
						}
					}
				}
			}

			free(buffer);
		}

		cxGetTokenPrivileges(hToken, token.privileges);

		CloseHandle(hToken);

		return ERROR_SUCCESS;
	}

	
}